# -*- coding: utf-8 -*-
"""
Created on Mon Jun 25 15:39:41 2012

@author: nzhao
"""

# Major library imports
from numpy import arange

# Enthought library imports
from enable.api import Component, ComponentEditor
from traits.api import Float, HasTraits, Int, Instance, List, Button
from traitsui.api import Item, Group, View

# Chaco imports
from chaco.api import create_line_plot, add_default_axes, add_default_grids, \
        OverlayPlotContainer, PlotLabel, create_scatter_plot, Legend, HPlotContainer
from chaco.tools.api import PanTool, ZoomTool, LegendTool, TraitsTool, DragZoom, LineInspector

class OverlappingPlotContainer(OverlayPlotContainer):
    """Simple container for creating a series of plots"""

    num_plots = Int
    data_list = List([])
    
    def __init__(self, *args, **kws):
        super(OverlayPlotContainer, self).__init__(*args, **kws)
        self.num_plots = len(self.data_list)
    
    def clear_data(self):
        self.data_list = []
        self.num_plots = len(self.data_list)
        len_comp = len(self.components)
        for i in range(len_comp):
            #print len_comp - i -1, ': ', self.components[len_comp - i -1]
            self.remove(self.components[len_comp - i -1])
        
    def add_data(self, data):
        self.data_list.append(data)
        self.num_plots = len(self.data_list)
        
    def plot_data(self, color_list):
        """Creates series of Bessel function plots"""
        plots = {}

        for i in range(self.num_plots):
            x, y = self.data_list[i]
            plot = create_line_plot((x, y), value_bounds= [-1.1, 1.1], color= color_list[i], width=2.0)
            if i == 0:
                value_mapper, index_mapper, legend = \
                    self._setup_plot_tools(plot)
            else:
                self._setup_mapper(plot, value_mapper, index_mapper)

            self.add(plot)
            plots["function_%d" % i] = plot

        # Set the list of plots on the legend
        legend.plots = plots

        # Add the traits inspector tool to the container
        self.tools.append(TraitsTool(self))
        
        
    def _setup_plot_tools(self, plot):
        """Sets up the background, and several tools on a plot"""
        # Make a white background with grids and axes
        plot.bgcolor = "white"
        add_default_grids(plot)
        add_default_axes(plot)

        # Allow white space around plot
        plot.index_range.tight_bounds = True
        plot.index_range.refresh()
        plot.value_range.tight_bounds = True
        plot.value_range.refresh()

        # The PanTool allows panning around the plot
        plot.tools.append(PanTool(plot))

        # The ZoomTool tool is stateful and allows drawing a zoom
        # box to select a zoom region.
        zoom = ZoomTool(plot, tool_mode="box", always_on=False)
        plot.overlays.append(zoom)

        # The DragZoom tool just zooms in and out as the user drags
        # the mouse vertically.
        dragzoom = DragZoom(plot, drag_button="right")
        plot.tools.append(dragzoom)

        # Add a legend in the upper right corner, and make it relocatable
        legend = Legend(component=plot, padding=10, align="ur")
        legend.tools.append(LegendTool(legend, drag_button="right"))
        plot.overlays.append(legend)
        plot.overlays.append(LineInspector(component=plot,
               axis="value", write_metadata=True, is_listener=True))
        plot.overlays.append(LineInspector(component=plot,
               axis="index", write_metadata=True, is_listener=True))

        return plot.value_mapper, plot.index_mapper, legend

    def _setup_mapper(self, plot, value_mapper, index_mapper):
        """Sets up a mapper for given plot"""
        plot.value_mapper = value_mapper
        value_mapper.range.add(plot.value)
        plot.index_mapper = index_mapper
        index_mapper.range.add(plot.index)

 
class PlotExample(HasTraits):
    size = (800, 700)
    title = "Simple Line Plot"  

    plot1 = Instance(Component)
    plot2 = Instance(Component)
    plotAll = Instance(HPlotContainer)
    do_plot = Button

    traits_view = View(
                    Group(
                        Item('plotAll', editor=ComponentEditor(size=size),
                             show_label=False),
                        Item('do_plot'),
                        orientation="vertical"),
                    resizable=True, title=title,
                    width=size[0], height=size[1]
                    )
    def _do_plot_fired(self):
        self.plot1.add_data([[1,2,3],[2,3,4]])
        self.plot1.add_data([[10,20,30],[20,30,40]])
        self.plot1.plot_data(['red','blue'])
        
        self.plot2.add_data([[3,4,5],[2,3,4]])
        self.plot2.add_data([[10,20,30],[25,35,45]])
        self.plot2.plot_data(['red','blue'])
        
        self.plotAll = HPlotContainer(self.plot1, self.plot2, stack_order="left_to_right")


    def _plot1_default(self):
        container = OverlappingPlotContainer(padding=50, fill_padding=True, bgcolor="lightgray", use_backbuffer=True)
        return container
    def _plot2_default(self):
        container = OverlappingPlotContainer(padding=50, fill_padding=True, bgcolor="lightgray", use_backbuffer=True)
        return container




if __name__ == "__main__":
    demo = PlotExample()
    demo.configure_traits()
    
